# -*- coding: utf-8 -*-
"""
Created on Thu Jun 21 16:21:05 2012

@author: nzhao
"""
from traits.api import HasTraits, Int

import numpy as np
from scipy.linalg.matfuncs import expm2
from numpy.linalg import matrix_power

from math_qm.Constant import X, Y, Z

Default_Pseudo_Field = [[0.,0.,0.],[0.,0.,0.]]

class CPMG(HasTraits):
    spin_num = Int(1)
    
    def __init__(self, spin_num = 1):
        self.spin_num = spin_num

    def hamiltonian(self, field_vec):
        res = 0.5 * field_vec[0] * X + 0.5 * field_vec[1] * Y + 0.5 * field_vec[2] * Z
        return res
    
    def rotation_vector(self, u_mat):
        Ct = 0.5 * np.trace( u_mat ).real
        Xt = 0.5 * np.trace( np.dot(u_mat, X) ).imag
        Yt = 0.5 * np.trace( np.dot(u_mat, Y) ).imag
        Zt = 0.5 * np.trace( np.dot(u_mat, Z) ).imag
        return [Ct, Xt, Yt, Zt]
    
    def coherence(self, timelist=[0.0], pseudo_field= Default_Pseudo_Field, pulse_num = 1):
        h0_vec, h1_vec= pseudo_field
        self.hami_0 = self.hamiltonian(h0_vec)
        self.hami_1 = self.hamiltonian(h1_vec)
        
        res = [self.coherenceAt(t, pseudo_field, pulse_num) for t in timelist]
        return np.array(res)

    def evolution_operators(self, pseudo_field = Default_Pseudo_Field, time = 0.0, pulse_num = 1):
        k = int(pulse_num/2)
        tau = time / 2 / pulse_num
        
        hami_A, hami_B = self.hami_0, self.hami_1

        uA = expm2(-1.j * tau * hami_A)
        uB = expm2(-1.j * tau * hami_B)
        uBuA = np.dot(uB, uA)
        uAuB = np.dot(uA, uB)
        
        if pulse_num == 1:
            uHahn_A = uBuA
            uHahn_B = uAuB 
            res = [uHahn_A, uHahn_B]
        else:# pulse_num > 1:
            uCP2_A  = np.dot(uAuB, uBuA)
            uCP2_B  = np.dot(uBuA, uAuB)
            uCP2k_A = matrix_power(uCP2_A, k)
            uCP2k_B = matrix_power(uCP2_B, k)
            
            if pulse_num % 2 == 0:
                res = [uCP2k_A, uCP2k_B]
            else:
                uCP2k1_A = np.dot(uBuA, uCP2k_A)
                uCP2k1_B = np.dot(uAuB, uCP2k_B)
                res = [uCP2k1_A, uCP2k1_B]
        return res
        
    def coherenceAt(self, time=0.0, pseudo_field= Default_Pseudo_Field, pulse_num = 1):
        evo_op_A, evo_op_B = self.evolution_operators(pseudo_field, time, pulse_num)
        evo_op_coherence = np.dot( np.conjugate(evo_op_B).T, evo_op_A)
        
        coherence_vect = self.rotation_vector(evo_op_coherence)
        left_vect      = self.rotation_vector(evo_op_B)
        right_vect     = self.rotation_vector(evo_op_A)
        
        return np.array([coherence_vect, left_vect, right_vect])
        
    def single_spinHalf_coherence(self, pseudo_field= Default_Pseudo_Field, time=0.0, pulse_num = 1):
        if pulse_num == 1:
            k = 0.5
        elif pulse_num  % 2 == 0:
            k = pulse_num / 2
        else:
            k = (pulse_num - 1 ) / 2 
        
        tau = time / 4 / k;

        h0_vec, h1_vec= pseudo_field
        
        h0     = np.linalg.norm(h0_vec)
        h1     = np.linalg.norm(h1_vec)
        
        h0_dir = np.array([0., 0., 1.])
        if h0 > 0.0: h0_dir = h0_vec / h0
        
        h1_dir = np.array([0., 0., 1.])
        if h1 > 0.0: h1_dir = h1_vec / h1

        dot01 = np.dot(h0_vec, h1_vec) / h0 / h1

        cross01_dir  = np.cross(h0_dir, h1_dir)

        phi0 = 2.0 * np.pi * h0 * tau
        phi1 = 2.0 * np.pi * h1 * tau

        R1 = np.cos(0.5 * phi0) * np.sin(0.5 * phi1) * h1_dir + np.cos(0.5 * phi1) * np.sin(0.5 * phi0) * h0_dir
        r1 = np.sin(0.5 * phi0) * np.sin(0.5 * phi1) * cross01_dir
        
        xi1  = np.array([0., 0., 0.])
        xi2k = np.array([0., 0., 0.])

        if pulse_num == 1:
            coherence = np.array([1 - 2.0 * np.linalg.norm(r1)**2, 0.0, 0.0])
            res = np.array([coherence, R1, r1, xi1, xi2k])
        elif pulse_num > 1:
            cosTheta1Half = np.cos(0.5 * phi0) * np.cos( 0.5 * phi1) - dot01 * np.sin(0.5 * phi0) * np.sin(0.5 * phi1)
            R2 = 2.0 * cosTheta1Half * R1
            r2 = 2.0 * np.cross(R1, r1)
            alpha = np.arctan(np.linalg.norm(r1) / cosTheta1Half)
            
            R2_dir = np.array([0., 0., 1.])
            r2_dir = np.array([0., 0., 1.])
            if np.linalg.norm(R2) >0.0: R2_dir = R2 / np.linalg.norm(R2)
            if np.linalg.norm(r2) >0.0: r2_dir = r2 / np.linalg.norm(r2)
            
            cosTheta2Half = np.cos(phi0) * np.cos(phi1) - dot01 * np.sin(phi0) * np.sin(phi1)
            theta2Half     = np.arccos(cosTheta2Half)
            sinTheta2kHalf = np.sin(k * theta2Half)
        
            R2k = sinTheta2kHalf * np.cos(alpha) * R2_dir
            r2k = sinTheta2kHalf * np.sin(alpha) * r2_dir
            
            if pulse_num % 2 == 0:            
                coherence = np.array([1 - 2.0 * np.linalg.norm(r2k)**2, 0.0, 0.0])
            else:
                xi1  = -r2 + 2.0 * cosTheta1Half * r1
                xi2k =  2.0 * (np.cross(r2k, R2k) - np.cos(k* theta2Half) * r2k )
                coherenceHahn = 1 - 2.0 * np.linalg.norm(r1)**2
                coherence2k   = 1 - 2.0 * np.linalg.norm(r2k)**2
                coherence     = np.array([coherenceHahn * coherence2k + np.dot(xi1, xi2k), 0.0, 0.0])
            
            res = np.array([coherence, R2k, r2k, xi1, xi2k])
        
        return res